import { MongoClient } from "mongodb";
import chalk  from 'chalk';

import { GovDocTextSplitter } from "./docParseHelper.js";
import {
  MONGO_URI,
  MONGO_DB,
  SOURCE_COLLECTION,
  DESTINATION_COLLECTION,
} from "./config.js";

async function transferData() {
  const client = new MongoClient(MONGO_URI, {
    useNewUrlParser: true,
    useUnifiedTopology: true,
  });

  try {
    // 连接到 MongoDB
    await client.connect();
    console.log(chalk.bold.yellow('mogodb connection connected'))
    const db = client.db(MONGO_DB);
    const sourceCollection = db.collection('documents_sp_rag');
    const destinationCollection = db.collection('documents_artical_fragments_lyk');

    // 创建一个游标来遍历集合
    const cursor = sourceCollection.find(
      //  {platform_id: { $ne: "rulaixiezuo" }}, // 过滤掉"如来写作" 仅查询以inner开头的
      // { platform_id: "inner_qstheory_hongqi" }, 
      // { platform_id: "inner_news_comments" }, 
      // {platform_id: "inner_cpcnews_all"},
      // { platform_id: "inner_people_innovativecase" },
      // {platform_id:"inner_qstheory"},
      // {platform_id:"inner_people"},
      { isCanDroped: false},
      { noCursorTimeout: true }, // 必须放在最后一行
    );

    // 批量处理数据
    const batchSize = 1000;
    let batch = [];

    for await (const doc of cursor) {
      const {
        _id,
        id,
        platform_id,
        category_type,
        origin_category,
        title,
        url,
        content,
      } = doc;
      const article_id = id;
      const document_id = _id.toString();
      //   如果小于 100 文章过短  忽略
      if (content.length < 100) continue;

      const cleanContent = content.replace(title, "").trim();
      const ariticleChunks = GovDocTextSplitter.splitText(cleanContent);

      let len = ariticleChunks.length;
      if (len === 0) continue;

      // 检查最后一个片段 如果太短 小于100的话 则与前一个合并
      if (len > 1 && ariticleChunks[len - 1].trim().length < 100) {
        ariticleChunks[len - 2] += ariticleChunks[len - 1];
        ariticleChunks.pop();
        len--;
      }
      console.log(chalk.green(`Processing article: ${url} -- ${title}`));
      for (let i = 0; i < len; i++) {
        const cleanedChunk = ariticleChunks[i].trim();
        const fragmentId = article_id + `_${i}`;
        const exists = await destinationCollection.findOne({ id: fragmentId });
        // 检查是否已存在
        if (exists) {
          console.log(`fragment with id=${fragmentId} already exists`);
          continue;
        }
        const fragment = {
          id: fragmentId,
          document_id: document_id,
          fragment_index: i,
          fragment_content: cleanedChunk,
          platform_id,
          category_type,
          origin_category,
          title,
          url,
        };
        batch.push(fragment);
      }
      // 批量插入 每次插入
      if (batch.length >= batchSize) {
        try {
          await destinationCollection.insertMany(batch, { ordered: false });
          console.log(chalk.bold.green(`Inserted batch of ${batch.length} documents`));
        } catch (err) {
          console.error(`Error inserting batch: ${err}`);
        }
        batch = [];
      }
    }

    // 插入剩余的文档 不够1000了 剩余的直接插入
    if (batch.length > 0) {
      try {
        await destinationCollection.insertMany(batch, { ordered: false });
        console.log(chalk.bold.green(`Inserted final batch of ${batch.length} documents`));
      } catch (err) {
        console.error(chalk.bold.yellow(`Error inserting final batch: ${err}`));
      }
      console.log(chalk.bold.yellow('success: all docs inserted'))
      
    }
  } catch (err) {
    console.error(err);
  } finally {
    console.log(chalk.bold.yellow('mogodb connection closed'))
    await client.close();
  }
}

// 执行数据转移
transferData().catch(console.error);
